#!/usr/bin/env python3
"""
测试数据生成脚本
生成各种类型的测试数据，包括漏洞、产品、攻击向量等
"""

import json
import random
import csv
import os
from datetime import datetime, timedelta
from typing import List, Dict
from faker import Faker

fake = Faker(['zh_CN', 'en_US'])

def create_test_directories():
    """创建测试数据目录"""
    directories = [
        'tests/data/raw',
        'tests/data/processed', 
        'tests/data/expected',
        'tests/data/mock_apis',
        'tests/reports'
    ]
    
    for directory in directories:
        os.makedirs(directory, exist_ok=True)
        print(f"Created directory: {directory}")

def generate_test_vulnerabilities(count=100) -> List[Dict]:
    """生成测试用漏洞数据"""
    vulnerabilities = []
    
    severities = ["LOW", "MEDIUM", "HIGH", "CRITICAL"]
    attack_vectors = ["NETWORK", "ADJACENT_NETWORK", "LOCAL", "PHYSICAL"]
    attack_complexity = ["LOW", "HIGH"]
    
    for i in range(count):
        year = random.randint(2020, 2024)
        cve_id = f"CVE-{year}-{10000 + i}"
        
        vuln = {
            "cve_id": cve_id,
            "title": fake.sentence(nb_words=6)[:-1],  # 去掉末尾的句号
            "description": fake.text(max_nb_chars=500),
            "severity": random.choice(severities),
            "cvss_score": round(random.uniform(1.0, 10.0), 1),
            "cvss_vector": f"CVSS:3.1/AV:{random.choice(['N', 'A', 'L', 'P'])}/AC:{random.choice(['L', 'H'])}/PR:{random.choice(['N', 'L', 'H'])}/UI:{random.choice(['N', 'R'])}/S:{random.choice(['U', 'C'])}/C:{random.choice(['N', 'L', 'H'])}/I:{random.choice(['N', 'L', 'H'])}/A:{random.choice(['N', 'L', 'H'])}",
            "attack_vector": random.choice(attack_vectors),
            "attack_complexity": random.choice(attack_complexity),
            "published_date": (datetime.now() - timedelta(days=random.randint(1, 1000))).isoformat(),
            "modified_date": (datetime.now() - timedelta(days=random.randint(1, 30))).isoformat(),
            "affected_products": [
                f"{fake.company()} {fake.word().title()} Server",
                f"{fake.word().title()} {random.choice(['Framework', 'Library', 'Application'])} {random.randint(1, 10)}.{random.randint(0, 9)}.{random.randint(0, 9)}"
            ],
            "references": [
                f"https://cve.mitre.org/cgi-bin/cvename.cgi?name={cve_id}",
                f"https://nvd.nist.gov/vuln/detail/{cve_id}",
                f"https://www.exploit-db.com/exploits/{random.randint(40000, 50000)}",
                fake.url()
            ],
            "cwe_ids": [f"CWE-{random.randint(1, 1000)}" for _ in range(random.randint(1, 3))],
            "tags": random.sample([
                "remote", "local", "privilege-escalation", "code-execution", 
                "sql-injection", "xss", "buffer-overflow", "denial-of-service",
                "authentication-bypass", "information-disclosure"
            ], k=random.randint(1, 4))
        }
        vulnerabilities.append(vuln)
    
    return vulnerabilities

def generate_test_products(count=50) -> List[Dict]:
    """生成测试用产品数据"""
    products = []
    
    vendors = ["Apache", "Microsoft", "Oracle", "Google", "IBM", "Cisco", "Adobe", "VMware"]
    product_types = ["Server", "Framework", "Library", "Database", "Browser", "OS", "Application"]
    
    for i in range(count):
        vendor = random.choice(vendors)
        product_type = random.choice(product_types)
        
        product = {
            "cpe_id": f"cpe:2.3:a:{vendor.lower()}:{fake.word()}:{random.randint(1, 10)}.{random.randint(0, 9)}.{random.randint(0, 9)}:*:*:*:*:*:*:*",
            "vendor": vendor,
            "product": f"{fake.word().title()} {product_type}",
            "version": f"{random.randint(1, 10)}.{random.randint(0, 9)}.{random.randint(0, 9)}",
            "update": random.choice(["", "update1", "patch2", "sp3"]),
            "edition": random.choice(["", "enterprise", "professional", "community"]),
            "language": random.choice(["*", "en", "zh-cn", "ja"]),
            "part": random.choice(["a", "h", "o"]),  # application, hardware, operating system
            "description": fake.text(max_nb_chars=200),
            "homepage": fake.url(),
            "download_url": fake.url(),
            "release_date": fake.date_between(start_date='-5y', end_date='today').isoformat(),
            "end_of_life": fake.date_between(start_date='today', end_date='+2y').isoformat() if random.random() > 0.7 else None
        }
        products.append(product)
    
    return products

def generate_test_attack_vectors(count=30) -> List[Dict]:
    """生成测试用攻击向量数据"""
    attack_vectors = []
    
    techniques = [
        "T1059", "T1105", "T1055", "T1083", "T1082", "T1016", "T1033", "T1007",
        "T1018", "T1057", "T1012", "T1124", "T1049", "T1065", "T1090", "T1071"
    ]
    
    tactics = [
        "Initial Access", "Execution", "Persistence", "Privilege Escalation",
        "Defense Evasion", "Credential Access", "Discovery", "Lateral Movement",
        "Collection", "Command and Control", "Exfiltration", "Impact"
    ]
    
    for i in range(count):
        technique_id = random.choice(techniques)
        
        attack_vector = {
            "technique_id": technique_id,
            "name": fake.sentence(nb_words=4)[:-1],
            "description": fake.text(max_nb_chars=300),
            "tactics": random.sample(tactics, k=random.randint(1, 3)),
            "platforms": random.sample(["Windows", "Linux", "macOS", "Network"], k=random.randint(1, 3)),
            "data_sources": random.sample([
                "Process monitoring", "File monitoring", "Network traffic",
                "Windows event logs", "Authentication logs", "API monitoring"
            ], k=random.randint(1, 4)),
            "detection": fake.text(max_nb_chars=200),
            "mitigation": fake.text(max_nb_chars=200),
            "permissions": random.sample(["User", "Administrator", "SYSTEM", "root"], k=random.randint(1, 2)),
            "effective_permissions": random.choice(["User", "Administrator", "SYSTEM", "root"]),
            "network_requirements": random.choice([True, False]),
            "remote": random.choice([True, False]),
            "difficulty": random.choice(["Low", "Medium", "High"])
        }
        attack_vectors.append(attack_vector)
    
    return attack_vectors

def generate_test_html_pages(count=20) -> None:
    """生成测试用HTML页面"""
    templates = [
        # CVE详情页面模板
        """
        <html>
        <head><title>{cve_id} - CVE Details</title></head>
        <body>
            <h1>{cve_id}</h1>
            <div class="description">{description}</div>
            <div class="severity">Severity: {severity}</div>
            <div class="cvss-score">CVSS Score: {cvss_score}</div>
            <div class="published-date">Published: {published_date}</div>
            <div class="modified-date">Modified: {modified_date}</div>
            <ul class="affected-products">
                {affected_products_html}
            </ul>
            <div class="references">
                {references_html}
            </div>
        </body>
        </html>
        """,
        
        # 安全公告页面模板
        """
        <html>
        <head><title>Security Advisory - {title}</title></head>
        <body>
            <div class="post">
                <h2 class="title">{title}</h2>
                <div class="date">{date}</div>
                <div class="content">{content}</div>
                <div class="tags">{tags}</div>
            </div>
        </body>
        </html>
        """
    ]
    
    vulnerabilities = generate_test_vulnerabilities(count)
    
    for i, vuln in enumerate(vulnerabilities[:count]):
        template = random.choice(templates)
        
        affected_products_html = '\n'.join([f"<li>{product}</li>" for product in vuln['affected_products']])
        references_html = '\n'.join([f'<a href="{ref}">{ref}</a><br/>' for ref in vuln['references']])
        
        html_content = template.format(
            cve_id=vuln['cve_id'],
            title=vuln['title'],
            description=vuln['description'],
            severity=vuln['severity'],
            cvss_score=vuln['cvss_score'],
            published_date=vuln['published_date'],
            modified_date=vuln['modified_date'],
            affected_products_html=affected_products_html,
            references_html=references_html,
            date=vuln['published_date'],
            content=vuln['description'],
            tags=', '.join(vuln.get('tags', []))
        )
        
        filename = f"tests/data/raw/test_page_{i+1}.html"
        with open(filename, 'w', encoding='utf-8') as f:
            f.write(html_content)

def generate_mock_api_responses():
    """生成模拟API响应数据"""
    # NVD API响应
    nvd_response = {
        "resultsPerPage": 20,
        "startIndex": 0,
        "totalResults": 1000,
        "format": "NVD_CVE",
        "version": "2.0",
        "timestamp": datetime.now().isoformat(),
        "vulnerabilities": []
    }
    
    vulnerabilities = generate_test_vulnerabilities(20)
    for vuln in vulnerabilities:
        nvd_vuln = {
            "cve": {
                "id": vuln['cve_id'],
                "sourceIdentifier": "nvd@nist.gov",
                "published": vuln['published_date'],
                "lastModified": vuln['modified_date'],
                "vulnStatus": "Analyzed",
                "descriptions": [
                    {
                        "lang": "en",
                        "value": vuln['description']
                    }
                ],
                "metrics": {
                    "cvssMetricV31": [
                        {
                            "source": "nvd@nist.gov",
                            "type": "Primary",
                            "cvssData": {
                                "version": "3.1",
                                "vectorString": vuln['cvss_vector'],
                                "baseScore": vuln['cvss_score'],
                                "baseSeverity": vuln['severity']
                            }
                        }
                    ]
                },
                "weaknesses": [
                    {
                        "source": "nvd@nist.gov",
                        "type": "Primary",
                        "description": [
                            {
                                "lang": "en",
                                "value": cwe_id
                            } for cwe_id in vuln.get('cwe_ids', [])
                        ]
                    }
                ],
                "configurations": [],
                "references": [
                    {
                        "url": ref,
                        "source": "nvd@nist.gov",
                        "tags": ["Third Party Advisory"]
                    } for ref in vuln['references']
                ]
            }
        }
        nvd_response["vulnerabilities"].append(nvd_vuln)
    
    with open('tests/data/mock_apis/nvd_response.json', 'w', encoding='utf-8') as f:
        json.dump(nvd_response, f, indent=2, ensure_ascii=False)
    
    # GitHub API响应
    github_response = {
        "total_count": 100,
        "incomplete_results": False,
        "items": []
    }
    
    for i in range(10):
        repo = {
            "id": 100000 + i,
            "name": f"security-{fake.word()}",
            "full_name": f"{fake.user_name()}/security-{fake.word()}",
            "description": fake.text(max_nb_chars=100),
            "html_url": fake.url(),
            "created_at": fake.date_time_between(start_date='-2y', end_date='now').isoformat(),
            "updated_at": fake.date_time_between(start_date='-30d', end_date='now').isoformat(),
            "language": random.choice(["Python", "JavaScript", "Go", "Java", "C++"]),
            "topics": random.sample([
                "security", "vulnerability", "exploit", "cybersecurity",
                "pentest", "hacking", "malware", "forensics"
            ], k=random.randint(2, 4)),
            "stargazers_count": random.randint(10, 1000),
            "watchers_count": random.randint(5, 500),
            "score": random.uniform(0.5, 1.0)
        }
        github_response["items"].append(repo)
    
    with open('tests/data/mock_apis/github_response.json', 'w', encoding='utf-8') as f:
        json.dump(github_response, f, indent=2, ensure_ascii=False)

def generate_expected_results():
    """生成预期的测试结果"""
    # 实体抽取预期结果
    expected_entities = {
        "entities": [
            {
                "type": "VULNERABILITY",
                "name": "CVE-2023-12345",
                "properties": {
                    "severity": "HIGH",
                    "cvss_score": 8.5,
                    "description": "Apache HTTP Server path traversal vulnerability"
                }
            },
            {
                "type": "PRODUCT", 
                "name": "Apache HTTP Server",
                "properties": {
                    "vendor": "Apache",
                    "version": "2.4.50"
                }
            }
        ],
        "relations": [
            {
                "source": "CVE-2023-12345",
                "target": "Apache HTTP Server",
                "relation_type": "AFFECTS",
                "confidence": 0.95
            }
        ]
    }
    
    with open('tests/data/expected/entity_extraction_result.json', 'w', encoding='utf-8') as f:
        json.dump(expected_entities, f, indent=2, ensure_ascii=False)
    
    # 图查询预期结果
    expected_graph_query = {
        "query": "MATCH (v:Vulnerability)-[:AFFECTS]->(p:Product) WHERE v.severity = 'HIGH' RETURN v.cve_id, p.name LIMIT 10",
        "results": [
            {
                "v.cve_id": "CVE-2023-12345",
                "p.name": "Apache HTTP Server"
            },
            {
                "v.cve_id": "CVE-2023-12346", 
                "p.name": "Nginx Web Server"
            }
        ],
        "metadata": {
            "execution_time": 0.85,
            "result_count": 2
        }
    }
    
    with open('tests/data/expected/graph_query_result.json', 'w', encoding='utf-8') as f:
        json.dump(expected_graph_query, f, indent=2, ensure_ascii=False)

def generate_performance_baseline():
    """生成性能基线数据"""
    baseline = {
        "api_endpoints": {
            "/api/vulnerabilities/search": {
                "avg_response_time": 120,  # ms
                "95th_percentile": 200,
                "max_response_time": 500,
                "requests_per_second": 100,
                "error_rate": 0.001
            },
            "/api/vulnerabilities/{id}": {
                "avg_response_time": 80,
                "95th_percentile": 150,
                "max_response_time": 300,
                "requests_per_second": 200,
                "error_rate": 0.0005
            },
            "/api/graph/query": {
                "avg_response_time": 800,
                "95th_percentile": 1200,
                "max_response_time": 2000,
                "requests_per_second": 20,
                "error_rate": 0.01
            }
        },
        "system_metrics": {
            "cpu_usage": 45,  # %
            "memory_usage": 60,  # %
            "disk_io": 30,  # %
            "network_io": 25  # %
        },
        "data_processing": {
            "crawler_pages_per_minute": 50,
            "entity_extraction_per_minute": 500,
            "graph_updates_per_minute": 100,
            "data_quality_score": 0.92
        }
    }
    
    with open('tests/data/expected/performance_baseline.json', 'w', encoding='utf-8') as f:
        json.dump(baseline, f, indent=2, ensure_ascii=False)

def export_to_csv():
    """导出测试数据到CSV格式"""
    vulnerabilities = generate_test_vulnerabilities(200)
    products = generate_test_products(100)
    attack_vectors = generate_test_attack_vectors(50)
    
    # 导出漏洞数据
    with open('tests/data/raw/vulnerabilities.csv', 'w', newline='', encoding='utf-8') as f:
        writer = csv.DictWriter(f, fieldnames=vulnerabilities[0].keys())
        writer.writeheader()
        for vuln in vulnerabilities:
            # 处理列表类型字段
            vuln_copy = vuln.copy()
            vuln_copy['affected_products'] = '; '.join(vuln['affected_products'])
            vuln_copy['references'] = '; '.join(vuln['references'])
            vuln_copy['cwe_ids'] = '; '.join(vuln.get('cwe_ids', []))
            vuln_copy['tags'] = '; '.join(vuln.get('tags', []))
            writer.writerow(vuln_copy)
    
    # 导出产品数据
    with open('tests/data/raw/products.csv', 'w', newline='', encoding='utf-8') as f:
        writer = csv.DictWriter(f, fieldnames=products[0].keys())
        writer.writeheader()
        writer.writerows(products)
    
    # 导出攻击向量数据
    with open('tests/data/raw/attack_vectors.csv', 'w', newline='', encoding='utf-8') as f:
        writer = csv.DictWriter(f, fieldnames=attack_vectors[0].keys())
        writer.writeheader()
        for av in attack_vectors:
            av_copy = av.copy()
            av_copy['tactics'] = '; '.join(av['tactics'])
            av_copy['platforms'] = '; '.join(av['platforms'])
            av_copy['data_sources'] = '; '.join(av['data_sources'])
            av_copy['permissions'] = '; '.join(av['permissions'])
            writer.writerow(av_copy)

def main():
    """主函数"""
    print("🎯 开始生成测试数据...")
    
    # 创建目录
    create_test_directories()
    
    # 生成基础测试数据
    print("1. 生成漏洞数据...")
    vulnerabilities = generate_test_vulnerabilities(100)
    with open('tests/data/raw/test_vulnerabilities.json', 'w', encoding='utf-8') as f:
        json.dump(vulnerabilities, f, indent=2, ensure_ascii=False)
    
    print("2. 生成产品数据...")
    products = generate_test_products(50)
    with open('tests/data/raw/test_products.json', 'w', encoding='utf-8') as f:
        json.dump(products, f, indent=2, ensure_ascii=False)
    
    print("3. 生成攻击向量数据...")
    attack_vectors = generate_test_attack_vectors(30)
    with open('tests/data/raw/test_attack_vectors.json', 'w', encoding='utf-8') as f:
        json.dump(attack_vectors, f, indent=2, ensure_ascii=False)
    
    # 生成HTML测试页面
    print("4. 生成HTML测试页面...")
    generate_test_html_pages(20)
    
    # 生成模拟API响应
    print("5. 生成模拟API响应...")
    generate_mock_api_responses()
    
    # 生成预期结果
    print("6. 生成预期测试结果...")
    generate_expected_results()
    
    # 生成性能基线
    print("7. 生成性能基线数据...")
    generate_performance_baseline()
    
    # 导出CSV格式
    print("8. 导出CSV格式数据...")
    export_to_csv()
    
    # 生成综合测试数据集
    comprehensive_dataset = {
        "metadata": {
            "generated_at": datetime.now().isoformat(),
            "version": "1.0.0",
            "description": "Comprehensive test dataset for vulnerability knowledge graph",
            "counts": {
                "vulnerabilities": len(vulnerabilities),
                "products": len(products),
                "attack_vectors": len(attack_vectors)
            }
        },
        "vulnerabilities": vulnerabilities,
        "products": products,
        "attack_vectors": attack_vectors
    }
    
    with open('tests/data/comprehensive_test_dataset.json', 'w', encoding='utf-8') as f:
        json.dump(comprehensive_dataset, f, indent=2, ensure_ascii=False)
    
    print("✅ 测试数据生成完成!")
    print(f"   - 漏洞数据: {len(vulnerabilities)} 条")
    print(f"   - 产品数据: {len(products)} 条")
    print(f"   - 攻击向量: {len(attack_vectors)} 条")
    print(f"   - HTML页面: 20 个")
    print(f"   - 文件位置: tests/data/")

if __name__ == "__main__":
    main() 